#!/usr/bin/env python3
"""
模型配置修复工具
为缺少config.json的模型目录生成配置文件
"""

import sys
import json
import argparse
from pathlib import Path

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.append(str(project_root))

from src.utils.logger import setup_logger


def create_config_json(model_dir: str, model_type: str = "bert", num_labels: int = None):
    """
    为模型目录创建config.json文件
    
    Args:
        model_dir: 模型目录路径
        model_type: 模型类型
        num_labels: 标签数量，如果不提供则从label_mappings.json推断
    """
    model_path = Path(model_dir)
    
    if not model_path.exists():
        raise ValueError(f"模型目录不存在: {model_dir}")
    
    # 检查是否已有config.json
    config_file = model_path / "config.json"
    if config_file.exists():
        print(f"config.json已存在: {config_file}")
        return
    
    # 尝试从label_mappings.json获取标签信息
    label_mappings_file = model_path / "label_mappings.json"
    id2label = {}
    label2id = {}
    
    if label_mappings_file.exists():
        print(f"从 {label_mappings_file} 读取标签映射")
        with open(label_mappings_file, 'r', encoding='utf-8') as f:
            mappings = json.load(f)
            
        if 'label_to_intent' in mappings:
            id2label = {str(k): v for k, v in mappings['label_to_intent'].items()}
        if 'intent_to_label' in mappings:
            label2id = mappings['intent_to_label']
        if 'num_classes' in mappings:
            num_labels = mappings['num_classes']
    
    # 如果仍然没有标签数量，使用默认值
    if num_labels is None:
        num_labels = len(id2label) if id2label else 8
    
    # 创建基本配置
    if model_type == "bert":
        config = {
            "_name_or_path": "bert-base-chinese",
            "architectures": ["BertForSequenceClassification"],
            "attention_probs_dropout_prob": 0.1,
            "classifier_dropout": None,
            "directionality": "bidi",
            "hidden_act": "gelu",
            "hidden_dropout_prob": 0.1,
            "hidden_size": 768,
            "id2label": id2label,
            "initializer_range": 0.02,
            "intermediate_size": 3072,
            "label2id": label2id,
            "layer_norm_eps": 1e-12,
            "max_position_embeddings": 512,
            "model_type": "bert",
            "num_attention_heads": 12,
            "num_hidden_layers": 12,
            "num_labels": num_labels,
            "pad_token_id": 0,
            "pooler_fc_size": 768,
            "pooler_num_attention_heads": 12,
            "pooler_num_fc_layers": 3,
            "pooler_size_per_head": 128,
            "pooler_type": "first_token_transform",
            "position_embedding_type": "absolute",
            "torch_dtype": "float32",
            "transformers_version": "4.54.0",
            "type_vocab_size": 2,
            "use_cache": True,
            "vocab_size": 21128
        }
    else:
        # 通用配置
        config = {
            "architectures": ["BertForSequenceClassification"],
            "model_type": model_type,
            "num_labels": num_labels,
            "id2label": id2label,
            "label2id": label2id,
            "hidden_size": 768,
            "vocab_size": 21128,
            "max_position_embeddings": 512,
            "type_vocab_size": 2,
            "hidden_dropout_prob": 0.1,
            "attention_probs_dropout_prob": 0.1,
        }
    
    # 保存配置
    with open(config_file, 'w', encoding='utf-8') as f:
        json.dump(config, f, ensure_ascii=False, indent=2)
    
    print(f"已创建 config.json: {config_file}")
    print(f"模型类型: {model_type}")
    print(f"标签数量: {num_labels}")
    if id2label:
        print(f"支持的意图: {list(id2label.values())}")


def validate_model_directory(model_dir: str):
    """
    验证模型目录的完整性
    
    Args:
        model_dir: 模型目录路径
    """
    model_path = Path(model_dir)
    
    if not model_path.exists():
        print(f"❌ 模型目录不存在: {model_dir}")
        return False
    
    print(f"🔍 检查模型目录: {model_dir}")
    
    # 必需文件列表
    required_files = {
        "config.json": "模型配置文件",
        "model.safetensors": "模型权重文件",
        "tokenizer.json": "分词器配置",
        "tokenizer_config.json": "分词器配置",
        "vocab.txt": "词汇表",
    }
    
    # 可选文件列表
    optional_files = {
        "label_mappings.json": "标签映射",
        "training_config.json": "训练配置",
        "training_history.json": "训练历史",
        "special_tokens_map.json": "特殊token映射",
        "training_args.bin": "训练参数"
    }
    
    all_good = True
    
    print("\n必需文件:")
    for filename, description in required_files.items():
        file_path = model_path / filename
        if file_path.exists():
            print(f"  ✅ {filename} - {description}")
        else:
            print(f"  ❌ {filename} - {description} (缺失)")
            all_good = False
    
    print("\n可选文件:")
    for filename, description in optional_files.items():
        file_path = model_path / filename
        if file_path.exists():
            print(f"  ✅ {filename} - {description}")
        else:
            print(f"  ⚠️  {filename} - {description} (缺失)")
    
    if all_good:
        print(f"\n🎉 模型目录完整，可以正常使用！")
    else:
        print(f"\n⚠️  模型目录缺少必需文件，可能无法正常加载")
    
    return all_good


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='模型配置修复工具')
    
    parser.add_argument(
        'model_dir',
        type=str,
        help='模型目录路径'
    )
    
    parser.add_argument(
        '--model-type',
        type=str,
        default='bert',
        help='模型类型 (默认: bert)'
    )
    
    parser.add_argument(
        '--num-labels',
        type=int,
        help='标签数量 (如果不提供则从label_mappings.json推断)'
    )
    
    parser.add_argument(
        '--validate-only',
        action='store_true',
        help='只验证模型目录，不创建配置文件'
    )
    
    parser.add_argument(
        '--force',
        action='store_true',
        help='强制覆盖已存在的config.json文件'
    )
    
    args = parser.parse_args()
    
    try:
        if args.validate_only:
            # 只验证目录
            validate_model_directory(args.model_dir)
        else:
            # 检查是否需要强制覆盖
            config_file = Path(args.model_dir) / "config.json"
            if config_file.exists() and not args.force:
                print(f"config.json已存在: {config_file}")
                print("使用 --force 参数强制覆盖")
                return
            
            # 创建配置文件
            create_config_json(
                args.model_dir,
                args.model_type,
                args.num_labels
            )
            
            # 验证结果
            print("\n验证修复结果:")
            validate_model_directory(args.model_dir)
            
    except Exception as e:
        print(f"❌ 错误: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
